import time

import streamlit as st

from api.chat.demo_api import Tongyi_llm

# 定义布局（侧边栏）
with st.sidebar:
    #  定义标题
    st.title("欢迎使用我的AI工具")

    st.divider()

    # 通义大模型密钥输入
    key = st.text_input("请输入通义大模型密钥：")

    # key获取跳转链接
    st.page_link(
        "https://bailian.console.aliyun.com/?spm=5176.28326591.0.0.40f76ee1vUR5Of&accounttraceid=d873c9fe82ea443d9f3170f30dd5690cfnng#/api_key",
        label="无key？请点我获取", icon="🌎")

    st.divider()

# 定义一个模拟流式输出的方法
def stream_data(result_text):
    for word in result_text.split(" "):
        yield word
        time.sleep(0.02)

# 定义一个聊天记录到会话管理中
if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "ai", "content": "你好,我是您的ai助手，请输入你的问题!"}]

print(st.session_state["messages"])
# 如果聊天记录会话中只有一个开场白，那我们需要将开场白用流式模拟的方式渲染出来
for message in st.session_state["messages"]:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# 定义聊天框界面
user_input = st.chat_input("请输入你的问题")
if user_input:
    st.session_state["messages"].append({"role": "human", "content": user_input})
    # 获取用户输入并调用模型对象，并处理其的历史消
    with st.chat_message("human"):
        st.write_stream(stream_data(" ".join(user_input)))

if user_input:
    if not key:
        st.info("请输入你的模型密钥")
        st.stop()
    tongyi = Tongyi_llm()
    tongyi.create_llm(key)
    with st.spinner("AI正在思考，请稍等"):
        message = tongyi.invoke(user_input)
        st.session_state["messages"].append({"role": "ai", "content": message})
        with st.chat_message("ai"):
            st.write_stream(stream_data(" ".join(message)))

print("测试")